from mask import (get_random_size_box_with_constraint, get_random_size_ellipse_with_constraint, bbox2mask, brush_stroke_mask, get_irregular_mask, random_bbox, bbox2mask_uncropping, random_cropping_bbox)
import cv2

import numpy as np
import os


def brush_stroke_mask_with_constraint(img_shape, mask, num_vertices=(4, 12), brush_width=(12, 40), max_loops=4):
    """
    在指定的 mask 区域内绘制自由形态的 brush stroke。
    
    Args:
        img_shape (tuple[int]): 图像大小 (H, W)。
        mask (numpy.ndarray): 现有的 mask，0 表示不可绘制区域，1 表示允许绘制区域。
        num_vertices (tuple[int]): 线条拐点范围。
        brush_width (tuple[int]): 线条宽度范围。
        max_loops (int): 最大线条数量。

    Returns:
        numpy.ndarray: 生成的 mask，形状 (H, W)。
    """
    img_h, img_w = img_shape
    brush_mask = np.zeros((img_h, img_w), dtype=np.uint8)  # 生成一个空白 mask

    loop_num = np.random.randint(1, max_loops + 1)

    for _ in range(loop_num):
        num_vertex = np.random.randint(num_vertices[0], num_vertices[1] + 1)
        width = np.random.randint(brush_width[0], brush_width[1] + 1)

        # 随机找到一个起点（必须在 mask=0 区域）
        yx = np.argwhere(mask == 0)
        if len(yx) == 0:  # 避免 mask 全是 0 的情况
            return mask
        start_y, start_x = yx[np.random.randint(0, len(yx))]

        # 生成随机角度
        angles = np.random.uniform(0, 2 * np.pi, num_vertex)
        vertices = [(start_x, start_y)]

        # 生成路径
        while len(vertices) <= 2:
            for angle in angles:
                r = np.random.randint(60, 80)  # 设定每个步长的随机半径
                new_x = int(np.clip(vertices[-1][0] + r * np.cos(angle), 0, img_w - 1))
                new_y = int(np.clip(vertices[-1][1] + r * np.sin(angle), 0, img_h - 1))

                if mask[new_y, new_x] == 0:  # 仅在 mask=0 的区域绘制
                    vertices.append((new_x, new_y))
                else:
                    break  # 如果超出可绘制区域，则终止路径
            angles = np.random.uniform(0, 2 * np.pi, num_vertex)

        # 在 brush_mask 上绘制
        if len(vertices) > 1:
            cv2.polylines(brush_mask, [np.array(vertices, np.int32)], isClosed=False, color=1, thickness=width)
            for v in vertices:
                cv2.circle(brush_mask, v, width // 2, 1, -1)

    # 仅在 mask=1 的区域上更新原 mask
    # mask[brush_mask == 1] = 1
    # return mask
    brush_mask[mask == 1] = 0
    return brush_mask



def get_inpaint_mask(image_size):
    # # mask未知区域为1， 已知区域为0
    max_box_shape = int(max(image_size)*0.3)
    # regular_mask = bbox2mask(image_size, random_bbox(img_shape=image_size, max_bbox_shape=(200, 200))) 
    regular_mask = bbox2mask(image_size, random_bbox(img_shape=image_size, max_bbox_shape=(max_box_shape, max_box_shape)))  
    irregular_mask = brush_stroke_mask(image_size, num_vertices=(4, 7),brush_width=(12, 30))
    mask = regular_mask | irregular_mask
    return mask



def get_irregular_mask(image_size, box_size):
    mask = bbox2mask(image_size, random_bbox(img_shape=image_size, max_bbox_shape=(box_size, box_size)))[:,:,0]
    irregular_mask = brush_stroke_mask_with_constraint(image_size, mask, num_vertices=(12, 15), brush_width=(100, 120),max_loops=1)
    return mask, irregular_mask



def get_irregular_mask_aware(image_size, mask_foreground):
    irregular_mask = brush_stroke_mask_with_constraint(image_size, mask_foreground, num_vertices=(15, 18), brush_width=(150, 180),max_loops=1)
    box_mask = get_random_size_box_with_constraint(image_size, mask_foreground, ratio=0.6, min_valid_ratio=0.6)
    ellipse_mask = get_random_size_ellipse_with_constraint(image_size, mask_foreground, ratio=0.6, min_valid_ratio=0.6)
    irregular_mask[mask_foreground == 1] = 0
    box_mask[mask_foreground == 1] = 0
    ellipse_mask[mask_foreground == 1] = 0

    return irregular_mask, box_mask, ellipse_mask

# for size in range(50, 901, 50):

#     mask, irregular_mask = get_irregular_mask((1024,1024), size)
#     print(mask.shape, irregular_mask.shape)
#     concatenated = np.hstack((mask, irregular_mask))
#     cv2.imwrite(f'irregular_mask/irregular_mask_{size}.png', concatenated*255)


abs_path = '/llmcapagroup1/test-bucket/xinyu/code/HunyuanDIT-PRE-main/hydit/data_loader/mask/'
mask_list = os.listdir(abs_path)
for mask_name in mask_list:
    mask = cv2.imread(abs_path + mask_name, cv2.IMREAD_GRAYSCALE) / 255
    mask[mask < 0.5] = 0
    mask[mask >= 0.5] = 1
    mask = 1 - mask
    print(mask.shape)
    irregular_mask, box_mask, ellipse_mask = get_irregular_mask_aware(mask.shape, mask)
    concatenated = np.hstack((mask, irregular_mask, box_mask, ellipse_mask))
    cv2.imwrite(f'irregular_mask/{mask_name}.png', concatenated*255)

